# Refer to MMYOLO
# Copyright (c) VCIP-NKU. All rights reserved.

from typing import Sequence, Union

import torch
import torch.nn as nn
from mmcv.cnn import ConvModule
from mmdet.utils import ConfigType, OptMultiConfig
from mmengine.model import BaseModule
from torch import Tensor

from mmyolo.registry import MODELS


@MODELS.register_module()
class SPPFBottleneck(BaseModule):
    """Spatial pyramid pooling - Fast (SPPF) layer for YOLOv5, YOLOX, and PPYOLOE.

    Args:
        in_channels (int): The input channels of this Module.
        out_channels (int): The output channels of this Module.
        kernel_sizes (int or tuple[int]): Sequential or number of kernel sizes of pooling layers. Defaults to 5.
        use_conv_first (bool): Whether to use conv before pooling layer.
            In YOLOv5 and YOLOX, this is set to True.
            In PPYOLOE, this is set to False. Defaults to True.
        mid_channels_scale (float): Channel multiplier, multiply `in_channels` by this amount to get `mid_channels`.
            This parameter is valid only when `use_conv_first=True`. Defaults to 0.5.
        conv_cfg (dict): Config dict for convolution layer. Defaults to None, which means using `Conv2d`.
        norm_cfg (dict): Config dict for normalization layer.
            Defaults to dict(type='BN', momentum=0.03, eps=0.001).
        act_cfg (dict): Config dict for activation layer.
            Defaults to dict(type='SiLU', inplace=True).
        init_cfg (dict or list[dict], optional): Initialization config dict. Defaults to None.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_sizes: Union[int, Sequence[int]] = 5,
        use_conv_first: bool = True,
        mid_channels_scale: float = 0.5,
        conv_cfg: ConfigType = None,
        norm_cfg: ConfigType = dict(type="BN", momentum=0.03, eps=0.001),
        act_cfg: ConfigType = dict(type="SiLU", inplace=True),
        init_cfg: OptMultiConfig = None,
    ):
        super().__init__(init_cfg)

        # Define the first convolution layer if `use_conv_first` is True
        if use_conv_first:
            mid_channels = int(in_channels * mid_channels_scale)
            self.conv1 = ConvModule(
                in_channels,
                mid_channels,
                1,
                stride=1,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=act_cfg,
            )
        else:
            mid_channels = in_channels
            self.conv1 = None

        self.kernel_sizes = kernel_sizes

        # Define pooling layers
        if isinstance(kernel_sizes, int):
            self.poolings = nn.MaxPool2d(
                kernel_size=kernel_sizes, stride=1, padding=kernel_sizes // 2
            )
            conv2_in_channels = mid_channels * 4
        else:
            self.poolings = nn.ModuleList(
                [
                    nn.MaxPool2d(kernel_size=ks, stride=1, padding=ks // 2)
                    for ks in kernel_sizes
                ]
            )
            conv2_in_channels = mid_channels * (len(kernel_sizes) + 1)

        # Define the second convolution layer
        self.conv2 = ConvModule(
            conv2_in_channels,
            out_channels,
            1,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg,
        )

    def forward(self, x: Tensor) -> Tensor:
        """Forward process.

        Args:
            x (Tensor): The input tensor.

        Returns:
            Tensor: The output tensor after applying SPPF.
        """
        if self.conv1:
            x = self.conv1(x)

        if isinstance(self.kernel_sizes, int):
            y1 = self.poolings(x)
            y2 = self.poolings(y1)
            x = torch.cat([x, y1, y2, self.poolings(y2)], dim=1)
        else:
            x = torch.cat([x] + [pooling(x) for pooling in self.poolings], dim=1)

        x = self.conv2(x)
        return x